package com.pro.framework.api.util;

import cn.hutool.json.JSONUtil;
import com.pro.framework.api.util.inner.SimpleRateLimiter;
import lombok.extern.slf4j.Slf4j;

import java.util.*;
import java.util.concurrent.*;
import java.util.function.BiFunction;
import java.util.function.Function;

@Slf4j
public class AsyncExecutorUtil {
    // 每个 key（如 openKeyId）对应一个限流器
    private static final Map<String, SimpleRateLimiter> limiterMap = new ConcurrentHashMap<>();

    /**
     * 获取限流器，若不存在则创建，单位：次/秒
     */
    public static SimpleRateLimiter getSimpleRateLimiter(String key, int permitsPerSecond) {
        return limiterMap.computeIfAbsent(key, k -> new SimpleRateLimiter(permitsPerSecond));
    }

    /**
     * 并发 执行
     */
    public static <T, R> List<R> executeConcurrently(List<T> inputList, Function<T, R> task, Executor executor) {
        List<CompletableFuture<R>> futures = inputList.stream()
                .map(i -> CompletableFuture.supplyAsync(() -> task.apply(i), executor))
                .toList();

        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

        return futures.stream()
                .map(CompletableFuture::join)
                .toList();
    }

    /**
     * 分批+并发 执行
     */
    public static <T, R> List<R> executeConcurrentlyPart(List<T> inputList, Function<T, R> task, Executor executor, Integer batchSize) {
        List<R> result = new ArrayList<>();
        List<List<T>> partitions = CollUtils.partitionList(inputList, batchSize);

        for (List<T> subList : partitions) {
            List<R> subResult = executeConcurrently(subList, task, executor); // 阻塞等待当前批完成
            result.addAll(subResult);
        }
        return result;
    }

    /**
     * 分批+并发执行（支持限流）
     */
    public static <T, R> List<R> executeConcurrentlyPart(
            List<T> inputList,
            Function<T, R> task,
            Executor executor,
            Integer batchSize,
            String rateLimiterKey,
            Integer permitsPerSecond) {

        List<R> result = new ArrayList<>();
        List<List<T>> partitions = CollUtils.partitionList(inputList, batchSize);

        // 获取限流器
        SimpleRateLimiter rateLimiter = null;
        if (rateLimiterKey != null && permitsPerSecond != null) {
            rateLimiter = AsyncExecutorUtil.getSimpleRateLimiter(rateLimiterKey, permitsPerSecond.intValue());
        }

        for (List<T> subList : partitions) {
            // 申请许可（如果配置了限流）
            if (rateLimiter != null) {
                rateLimiter.acquire();
            }

            List<R> subResult = executeConcurrently(subList, task, executor);
            result.addAll(subResult);
        }
        return result;
    }

    /**
     * 分批+并发执行（增强版）- 支持限流、重试、异常处理
     */
    public static <T, R> List<R> executeConcurrentlyPart(
            List<T> inputList,
            Function<T, R> task,
            Executor executor,
            Integer batchSize,
            String rateLimiterKey,
            Integer permitsPerSecond,
            // 新增参数
            Integer maxRetries) {
        return executeConcurrentlyPart(inputList, task, executor, batchSize, rateLimiterKey, permitsPerSecond, maxRetries, (o, e) -> {
            log.warn("{} 重试{}次 还是失败 {}", rateLimiterKey, maxRetries, JSONUtil.toJsonStr(o));
            return null;
        }).stream().filter(Objects::nonNull).toList();
    }

    public static <T, R> List<R> executeConcurrentlyPart(
            List<T> inputList,
            Function<T, R> task,
            Executor executor,
            Integer batchSize,
            String rateLimiterKey,
            Integer permitsPerSecond,
            // 新增参数
            Integer maxRetries,
            BiFunction<T, Exception, R> exceptionHandler) {

        List<R> result = new ArrayList<>();
        List<List<T>> partitions = CollUtils.partitionList(inputList, batchSize);

        // 获取限流器
        SimpleRateLimiter rateLimiter = null;
        if (rateLimiterKey != null && permitsPerSecond != null) {
            rateLimiter = AsyncExecutorUtil.getSimpleRateLimiter(rateLimiterKey, permitsPerSecond.intValue());
        }

        for (List<T> subList : partitions) {
            // 申请许可（如果配置了限流）
            if (rateLimiter != null) {
                rateLimiter.acquire();
            }

            // 执行当前批次（支持重试）
            List<R> subResult = executeBatchWithRetry(subList, task, executor, maxRetries, exceptionHandler);
            result.addAll(subResult);
        }
        return result;
    }

    /**
     * 轻量级部分重试版本
     * - 只重试失败项；
     * - 并发控制交给外部 Executor；
     * - 简化异常处理；
     */
    public static <T, R> List<R> executeBatchWithRetry(
            List<T> batch,
            Function<T, R> task,
            Executor executor,
            int maxRetries,
            BiFunction<T, Exception, R> exceptionHandler) {

        if (batch == null || batch.isEmpty()) {
            return Collections.emptyList();
        }

        List<R> results = new ArrayList<>(batch.size());
        List<T> remaining = new ArrayList<>(batch);

        for (int attempt = 0; attempt <= maxRetries && !remaining.isEmpty(); attempt++) {
            Map<T, CompletableFuture<R>> futureMap = new LinkedHashMap<>();
            List<T> failed = new ArrayList<>();

            for (T item : remaining) {
                futureMap.put(item, CompletableFuture.supplyAsync(() -> {
                    try {
                        return task.apply(item);
                    } catch (Exception e) {
                        // 关键修改：包装异常，让外层判断是否要降级
                        throw new CompletionException(e);
                    }
                }, executor));
            }

            CompletableFuture.allOf(futureMap.values().toArray(new CompletableFuture[0])).join();

            // 处理每个任务的结果
            for (Map.Entry<T, CompletableFuture<R>> entry : futureMap.entrySet()) {
                T item = entry.getKey();
                try {
                    R result = entry.getValue().join();
                    results.add(result);
                } catch (Exception e) {
                    Throwable realCause = e instanceof CompletionException ? e.getCause() : e;

                    if (realCause instanceof TimeoutException) {
                        // 对于超时错误进行 sleep 或指数回退
                        if (attempt < maxRetries) {
                            try {
                                // 使用指数回退（例如：2^attempt * 100ms）
                                //noinspection BusyWait
                                Thread.sleep((long) Math.pow(2, attempt) * 100);
                            } catch (InterruptedException interruptedException) {
                                Thread.currentThread().interrupt();
                            }
                        }
                    }
                    // 关键修改：如果是最后一次尝试，使用异常处理器
                    if (attempt >= maxRetries) {
                        if (exceptionHandler != null) {
                            results.add(exceptionHandler.apply(item, (Exception) realCause));
                        }
                    } else {
                        failed.add(item);
                    }
                }
            }

            remaining = failed;
        }

        return results;
    }
    /**
     * 并发执行 + 立即降级
     * 每个任务失败后立即调用降级处理器，不重试
     */
    public static <T, R> List<R> executeWithImmediateDegrade(
            List<T> inputList,
            Function<T, R> task,
            Executor executor,
            BiFunction<T, Exception, R> degradeHandler) {

        if (inputList == null || inputList.isEmpty()) {
            return Collections.emptyList();
        }

        // 为每个任务创建 CompletableFuture
        List<CompletableFuture<R>> futures = inputList.stream()
                .map(item -> CompletableFuture.supplyAsync(() -> {
                    try {
                        // 执行主要任务
                        return task.apply(item);
                    } catch (Exception e) {
                        // 关键：立即调用降级处理器
                        log.debug("任务执行失败，立即降级处理: {}", item, e);
                        return degradeHandler.apply(item, e);
                    }
                }, executor))
                .toList();

        // 等待所有任务完成
        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();

        // 收集结果
        return futures.stream()
                .map(CompletableFuture::join)
                .toList();
    }
    /**
     * 处理批次结果
     */
    private static <T, R> List<R> processBatchResults(
            List<T> batch,
            List<CompletableFuture<R>> futures,
            BiFunction<T, Exception, R> exceptionHandler) {

        List<R> results = new ArrayList<>();
        for (int i = 0; i < futures.size(); i++) {
            T item = batch.get(i);
            CompletableFuture<R> future = futures.get(i);

            try {
                R result = future.join();
                results.add(result);
            } catch (Exception e) {
                if (exceptionHandler != null) {
                    // 使用异常处理器处理单个失败任务
                    R fallbackResult = exceptionHandler.apply(item, e);
                    results.add(fallbackResult);
                } else {
                    // 没有异常处理器，直接抛出
                    throw new RuntimeException("任务执行失败: " + item, e);
                }
            }
        }
        return results;
    }

    /**
     * 处理整个批次失败的情况
     */
    private static <T, R> List<R> handleFailedBatch(
            List<T> batch,
            Exception batchException,
            BiFunction<T, Exception, R> exceptionHandler) {

        if (exceptionHandler != null) {
            return batch.stream()
                    .map(item -> exceptionHandler.apply(item, batchException))
                    .toList();
        }
        throw new RuntimeException("批次执行失败", batchException);
    }
}
